import pyopencl as cl
import numpy as np
from openclinfo import *

class Potential():
    def __init__(self,context,flags=''):
        self.potentials = []
        cl_code = file('potential.cl','r').read()
        self.prg = cl.Program(context,cl_code).build(options=flags)
        #

    def add_potential(self,pot):
        self.potentials.append(pot)

    def cal_f_e_slow(self,que,cfg):
        #glb_s = (cfg.n_atoms,)
        sglb = (GLOBAL_SIZE,)
        sloc = (LOCAL_SIZE,)
        self.prg.zero_f_e_vir(que,sglb,sloc,np.int32(cfg.n_atoms),cfg._f_e,cfg._vir)
        for pot in self.potentials:
            pot.add_f_e_slow(que,sglb,sloc,cfg)

    def cal_f_e_vir_all(self,que,cfg,nei):
        #glb_s = (cfg.n_atoms,)
        sglb = (GLOBAL_SIZE,)
        sloc = (LOCAL_SIZE,)
        self.prg.zero_f_e_vir(que,sglb,sloc,np.int32(cfg.n_atoms),cfg._f_e,cfg._vir)
        for pot in self.potentials:
            pot.add_f_e_vir_all(que,sglb,sloc,cfg,nei)
